package main

import (
	"fmt"
	"math"
	"strings"
)

var M = math.Inf(1)

func main() {
	matrix := [][]float64{
		//s  t  y  x  z
		{0, 6, 7, M, M},  //s
		{M, 0, 8, 5, -4}, //t 最初 t->y:8  更改 t-y:-8
		{M, M, 0, -3, 9}, //y //最初 y->s:M  新增 y->s:1
		{M, -2, M, 0, M}, //x
		{2, M, M, 7, 0},  //z
	}
	startNode, endNode := 0, 4
	loop, distance, path := bellmanFord2(matrix, startNode)
	fmt.Println("是否有负权环？", loop)
	fmt.Printf("%s->%s最短距离：%.0f，路线图：%v\n", nodeNames[startNode], nodeNames[endNode], distance[endNode], PathLineName(getPathLine(startNode, endNode, path)))

}

var nodeNames = map[int]string{
	0: "s",
	1: "t",
	2: "y",
	3: "x",
	4: "z",
}

type PathLineName []int

func (p PathLineName) String() string {
	sb := new(strings.Builder)
	for i, node := range p {
		if i > 0 {
			sb.WriteString("->")
		}
		sb.WriteString(nodeNames[node])
	}
	return sb.String()
}

func getPathLine(startNode, endNode int, path []int) []int {
	var pathNode = make([]int, len(path))
	prev := len(path) - 1
	pathNode[prev] = endNode
	for pri := path[endNode]; pri != -1 && pri != endNode; pri = path[pri] {
		prev--
		pathNode[prev] = pri
	}
	if prev > 0 && startNode != pathNode[prev] {
		prev--
		pathNode[prev] = startNode
	}
	pathNode = pathNode[prev:]
	return pathNode
}

func bellmanFord(matrix [][]float64, startNode int) (bool, []float64, []int) {
	n := len(matrix)
	path := make([]int, n)
	distance := make([]float64, n)
	for i := 0; i < n; i++ {
		path[i] = -1
		distance[i] = M
	}
	distance[startNode] = 0

	for k := 0; k < n-1; k++ {
		fmt.Printf("\n第%d轮操作\n", k+1)
		updated := false

		for i := 0; i < n; i++ {
			for j := 0; j < n; j++ {
				if matrix[i][j] != M && matrix[i][j]+distance[i] < distance[j] {
					updated = true
					distance[j] = matrix[i][j] + distance[i]
					path[j] = i
					fmt.Printf("%s->%s:%v\n", nodeNames[i], nodeNames[j], distance[j])
				}
			}
		}
		if !updated {
			break
		}
	}

	for i := 0; i < n; i++ {
		for j := 0; j < n; j++ {
			if matrix[i][j] != M && matrix[i][j]+distance[i] < distance[j] {
				return true, distance, path
			}
		}
	}

	return false, distance, path
}

// 队列优化
func bellmanFord2(matrix [][]float64, startNode int) (bool, []float64, []int) {
	n := len(matrix)
	path := make([]int, n)
	distance := make([]float64, n)
	queue := make([]int, n)
	for i := 0; i < n; i++ {
		path[i] = -1
		distance[i] = M
		queue[i] = i
	}
	distance[startNode] = 0

	for len(queue) != 0 {
		i := queue[0]
		queue = queue[1:]
		for j := 0; j < n; j++ {
			if matrix[i][j] != M && matrix[i][j]+distance[i] < distance[j] {
				distance[j] = matrix[i][j] + distance[i]
				path[j] = i
				queue = append(queue, j)
				fmt.Printf("%s->%s:%v\n", nodeNames[i], nodeNames[j], distance[j])
			}
		}
	}

	for i := 0; i < n; i++ {
		for j := 0; j < n; j++ {
			if matrix[i][j] != M && matrix[i][j]+distance[i] < distance[j] {
				return true, distance, path
			}
		}
	}

	return false, distance, path
}
